{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import os.path\n",
    "import random\n",
    "import numpy as np\n",
    "import numpy.random\n",
    "import pdb\n",
    "import random as rn\n",
    "import os.path\n",
    "import scipy.misc\n",
    "import time\n",
    "import tensorflow as tf\n",
    "import io\n",
    "print(tf.__version__)\n",
    "\n",
    "import srez_input_sia\n",
    "import srez_model_sia\n",
    "\n",
    "FLAGS = tf.app.flags.FLAGS\n",
    "\n",
    "# Configuration (alphabetically)\n",
    "tf.app.flags.DEFINE_integer('batch_size', 64, \"Number of samples per batch.\")\n",
    "tf.app.flags.DEFINE_string('checkpoint_dir', 'checkpoint', \"Output folder where checkpoints are dumped.\")\n",
    "tf.app.flags.DEFINE_integer('checkpoint_period', 1000, \"Number of batches in between checkpoints\")\n",
    "tf.app.flags.DEFINE_string('dataset', 'dataset', \"Path to the dataset directory.\")\n",
    "tf.app.flags.DEFINE_float('epsilon', 1e-8, \"Fuzz term to avoid numerical instability\")\n",
    "tf.app.flags.DEFINE_float('gene_l1_factor', .90, \"Multiplier for generator L1 loss term\")\n",
    "tf.app.flags.DEFINE_float('learning_beta1', 0.5, \"Beta1 parameter used for AdamOptimizer\")\n",
    "tf.app.flags.DEFINE_float('learning_rate_start', 0.00020, \"Starting learning rate used for AdamOptimizer\")\n",
    "tf.app.flags.DEFINE_integer('learning_rate_half_life', 5000, \"Number of batches until learning rate is halved\")\n",
    "tf.app.flags.DEFINE_bool('log_device_placement', False, \"Log the device where variables are placed.\")\n",
    "tf.app.flags.DEFINE_integer('sample_size', 32, \"Image sample size in pixels. \")\n",
    "tf.app.flags.DEFINE_integer('summary_period', 1000,\"Number of batches between summary data dumps\")\n",
    "tf.app.flags.DEFINE_integer('random_seed', 0, \"Seed used to initialize rng.\")\n",
    "tf.app.flags.DEFINE_integer('init_layer_size', 512, \"Initial #channels of the first layer in the generator.\")\n",
    "tf.app.flags.DEFINE_integer('test_vectors', 16,  \"\"\"Number of features to use for testing\"\"\")      \n",
    "tf.app.flags.DEFINE_string('train_dir', 'train', \"Output folder where the synthesized images during training are dumped.\")                      \n",
    "tf.app.flags.DEFINE_string('log_dir', 'logs/', \"Output folder where training logs are dumped.\")      \n",
    "tf.app.flags.DEFINE_string('training_img_dir', '../CASIA/CASIA-WebFace/', \"Output folder where training logs are dumped.\")\n",
    "tf.app.flags.DEFINE_string('testing_img_dir', '../CASIA/CASIA-WebFace/val.txt', \"Output folder where training logs are dumped.\")\n",
    "tf.app.flags.DEFINE_integer('train_time', 10000,  \"Time in minutes to train the model\")\n",
    "tf.app.flags.DEFINE_integer('HR_size', 32,  \"Time in minutes to train the model\")\n",
    "tf.app.flags.DEFINE_string('training_list', 'pairwise.txt', \"Pairwise file list for training set [path1 path2 same_or_not]\")      \n",
    "tf.app.flags.DEFINE_string('validation_list', 'pairwise-val2.txt', \"Pairwise file list for validation set\")      \n",
    "\n",
    "tf.app.flags.DEFINE_bool('useRecTerm', True,  \"Set whether the reconstruction term is used or not\")\n",
    "\n",
    "label=[]\n",
    "\n",
    "def variable_summaries(var):\n",
    "  \"\"\"Attach a lot of summaries to a Tensor (for TensorBoard visualization).\"\"\"\n",
    "  with tf.name_scope('summaries'):\n",
    "    mean = tf.reduce_mean(var)\n",
    "    tf.summary.scalar('mean', mean)\n",
    "    with tf.name_scope('stddev'):\n",
    "      stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))\n",
    "    tf.summary.scalar('stddev', stddev)\n",
    "    tf.summary.scalar('max', tf.reduce_max(var))\n",
    "    tf.summary.scalar('min', tf.reduce_min(var))\n",
    "    tf.summary.histogram('histogram', var)\n",
    "    \n",
    "\n",
    "def prepare_dirs(delete_train_dir=False):\n",
    "    # Create checkpoint dir (do not delete anything)\n",
    "    if not tf.gfile.Exists(FLAGS.checkpoint_dir):\n",
    "        tf.gfile.MakeDirs(FLAGS.checkpoint_dir)\n",
    "    \n",
    "    # Cleanup train dir\n",
    "    if delete_train_dir:\n",
    "        if tf.gfile.Exists(FLAGS.train_dir):\n",
    "            tf.gfile.DeleteRecursively(FLAGS.train_dir)\n",
    "        if not tf.gfile.Exists(FLAGS.train_dir):\n",
    "            tf.gfile.MakeDirs(FLAGS.train_dir)\n",
    "\n",
    "    fn=FLAGS.training_img_dir\n",
    "    fn2=FLAGS.testing_img_dir\n",
    "\n",
    "\n",
    "    #~ filenames = tf.gfile.ListDirectory(FLAGS.dataset)\n",
    "    #~ filenames = sorted(filenames)\n",
    "    #~ random.shuffle(filenames)\n",
    "    #~ filenames = [os.path.join(FLAGS.dataset, f) for f in filenames]\n",
    "\n",
    "    return FLAGS.training_img_dir\n",
    "\n",
    "\n",
    "def setup_tensorflow():\n",
    "    # Create session\n",
    "    config = tf.ConfigProto(log_device_placement=FLAGS.log_device_placement)\n",
    "    config.gpu_options.allow_growth = True\n",
    "    sess = tf.Session(config=config)\n",
    "\n",
    "    # Initialize rng with a deterministic seed\n",
    "    with sess.graph.as_default():\n",
    "        tf.set_random_seed(FLAGS.random_seed)\n",
    "        \n",
    "    random.seed(FLAGS.random_seed)\n",
    "    np.random.seed(FLAGS.random_seed)\n",
    "\n",
    "    summary_writer = tf.summary.FileWriter(FLAGS.log_dir, sess.graph)\n",
    "\n",
    "    return sess, summary_writer\n",
    "\n",
    "def _summarize_progress(sess, feature, label, gene_output, batch, suffix, max_samples=17):\n",
    "\n",
    "    size = [label.shape[1], label.shape[2]]\n",
    "\n",
    "    nearest = tf.image.resize_nearest_neighbor(feature, size)\n",
    "    nearest = tf.maximum(tf.minimum(nearest, 1.0), 0.0)\n",
    "\n",
    "    bicubic = tf.image.resize_bicubic(feature, size)\n",
    "    bicubic = tf.maximum(tf.minimum(bicubic, 1.0), 0.0)\n",
    "\n",
    "    clipped = tf.maximum(tf.minimum(gene_output, 1.0), 0.0)\n",
    "    #~ pdb.set_trace()\n",
    "    image   = tf.concat([nearest, bicubic, clipped, label], 2)\n",
    "    tf.summary.image(\"myResults\", image)\n",
    "\n",
    "    imagex = image[0:max_samples,:,:,:]\n",
    "    image = tf.concat([imagex[i,:,:,:] for i in range(max_samples)], 0)\n",
    "    image = sess.run(image)\n",
    "        \n",
    "\n",
    "    filename = 'batch%06d_%s.png' % (batch, suffix)\n",
    "    filename = os.path.join(FLAGS.train_dir, filename)\n",
    "    scipy.misc.toimage(image, cmin=0., cmax=1.).save(filename)\n",
    "    print(\"    Saved %s\" % (filename,))\n",
    "    \n",
    "    \n",
    "\n",
    "    return imagex\n",
    "\n",
    "def _save_checkpoint(sess, batch, saver):\n",
    "\n",
    "    oldname = 'checkpoint_old.txt'\n",
    "    newname = 'checkpoint_new.txt'\n",
    "    oldname = os.path.join(FLAGS.checkpoint_dir, oldname)\n",
    "    newname = os.path.join(FLAGS.checkpoint_dir, newname)\n",
    "\n",
    "    # Delete oldest checkpoint\n",
    "    try:\n",
    "        tf.gfile.Remove(oldname)\n",
    "        tf.gfile.Remove(oldname + '.meta')\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "    # Rename old checkpoint\n",
    "    try:\n",
    "        tf.gfile.Rename(newname, oldname)\n",
    "        tf.gfile.Rename(newname + '.meta', oldname + '.meta')\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "    # Generate new checkpoint\n",
    "    \n",
    "    saver.save(sess, newname)\n",
    "\n",
    "    print(\"    Checkpoint saved\")\n",
    "\n",
    "\n",
    "def _train():\n",
    "    # Setup global tensorflow state\n",
    "    sess, summary_writer = setup_tensorflow()\n",
    "\n",
    "    # Prepare directories\n",
    "    filenames = prepare_dirs(delete_train_dir=True)\n",
    "  \n",
    "    \n",
    "    train_features, train_features2, train_labels, train_labels2, ylab, fn1, fn2, glen = \\\n",
    "    srez_input_sia.setup_inputs(sess, FLAGS.training_img_dir+FLAGS.training_list, \\\n",
    "                                image_size=FLAGS.HR_size, crop_size=128)\n",
    "    test_features, test_features2, test_labels, test_labels2, ytlab, fnt1, fnt2, tlen = \\\n",
    "    srez_input_sia.setup_inputs(sess, FLAGS.training_img_dir+FLAGS.validation_list, \\\n",
    "                                image_size=FLAGS.HR_size, crop_size=128, isTest=True)\n",
    "    print(\"Now we have %d training image pairs to be loaded...\"%(glen))\n",
    "\n",
    "   \n",
    "    # Add some noise during training (think denoising autoencoders)\n",
    "    noise_level = .03\n",
    "    noisy_train_features = train_features + tf.random_normal(train_features.get_shape(), stddev=noise_level)\n",
    "    noisy_train_features2 = train_features2 + tf.random_normal(train_features2.get_shape(), stddev=noise_level)\n",
    "\n",
    "    # Create and initialize model\n",
    "    #==========Network1==================\n",
    "    [gene_minput, gene_moutput,\n",
    "     gene_output, gene_var_list,\n",
    "    disc_real_output, disc_fake_output, disc_var_list, \n",
    "    gene_minput2, gene_moutput2,\n",
    "     gene_output2, gene_var_list2,\n",
    "    disc_real_output2, disc_fake_output2, disc_var_list2,\n",
    "    feat1, feat2] = srez_model_sia.create_model(sess, noisy_train_features, train_labels, noisy_train_features2, train_labels2, True)\n",
    "\n",
    "    with tf.name_scope('Generator_loss'):\n",
    "         gene_loss = srez_model_sia.create_generator_loss(disc_fake_output, gene_output, train_features)\n",
    "         tf.summary.scalar(\"Generator_loss\", gene_loss)\n",
    "    \n",
    "    disc_real_loss, disc_fake_loss =   srez_model_sia.create_discriminator_loss(disc_real_output, disc_fake_output)\n",
    "    with tf.name_scope('Discriminator_loss'):\n",
    "        disc_loss = tf.add(disc_real_loss, disc_fake_loss, name='disc_loss')\n",
    "        tf.summary.scalar(\"D_Real_loss\", disc_real_loss)\n",
    "        tf.summary.scalar(\"D_Fake_loss\", disc_fake_loss)\n",
    "        tf.summary.scalar(\"Discriminator_loss\", disc_loss)\n",
    "    \n",
    "    #============Siamese contrastive loss==================\n",
    "    margin = 0.5\n",
    "    labels_t = tf.cast(ylab, tf.float32)\n",
    "    labels_f = tf.cast(1-ylab, tf.float32)         # labels_ = !labels;\n",
    "    eucd2 = tf.pow(feat1- feat2, 2.0)              # Feature distance.  \n",
    "    eucd2 = tf.reduce_sum(eucd2, [1])\n",
    "#     eucd2 = tf.pow(gene_output- gene_output2, 2.0)  # Pixel distance. uncomment if use it. \n",
    "#     eucd2 = tf.reduce_sum(eucd2, [1,2,3])\n",
    "    eucd = tf.sqrt(eucd2+1e-10, name=\"eucd\")\n",
    "    C = tf.constant(margin, name=\"C\")\n",
    "    pos = labels_t * eucd2\n",
    "    neg = labels_f *tf.pow(tf.maximum(C- eucd, 0), 2)\n",
    "    losses = pos + neg\n",
    "    with tf.name_scope('Contrastive_loss'):\n",
    "        sialoss = tf.reduce_mean(losses, name=\"Contrastive_loss\")\n",
    "        tf.summary.scalar(\"Contrastive_loss\", sialoss)\n",
    "    \n",
    "    #~ gene_loss = tf.add(gene_loss, sialoss*0.1)\n",
    "    \n",
    "    (global_step, learning_rate, gene_minimize, disc_minimize, sia_minimize, disc_var_list) =  \\\n",
    "    srez_model_sia.create_optimizers(gene_loss, gene_var_list, disc_loss, disc_var_list, sialoss,gene_var_list)\n",
    "\n",
    "    # Train model\n",
    "\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    saver = tf.train.Saver()\n",
    "    summaries = tf.summary.merge_all()\n",
    "\n",
    "    lrval       = FLAGS.learning_rate_start\n",
    "    start_time  = time.time()\n",
    "    done  = False\n",
    "    batch = 0\n",
    "    #~ pdb.set_trace()\n",
    "    assert FLAGS.learning_rate_half_life % 10 == 0\n",
    "\n",
    "    # Cache test features and labels (they are small)\n",
    "    test_feature, test_label = sess.run([test_features, test_labels])\n",
    "    ops =  [disc_minimize, disc_real_loss, disc_fake_loss, gene_minimize, gene_loss, sialoss, sia_minimize]\n",
    "    critic_itrs=4\n",
    "    cri = glen/100/FLAGS.batch_size\n",
    "    batch_total = glen/FLAGS.batch_size\n",
    "    gene_loss2 = disc_real_loss2 = disc_fake_loss2= sia_loss2 = -1.234\n",
    "    \n",
    "    while not done:\n",
    "        batch += 1\n",
    "        feed_dict = {learning_rate : lrval, gene_minput: test_feature, gene_minput2: test_feature}       \n",
    "        _, disc_real_loss2, disc_fake_loss2,_, gene_loss2, sia_loss2, _= sess.run(ops, feed_dict=feed_dict)\n",
    "   \n",
    "        if batch % 50 == 0:\n",
    "            # Show we are alive\n",
    "            elapsed = int(time.time() - start_time)/60\n",
    "            print('Batch[%d/%d(%3.3f%%)], Epoch[%2d] G_Loss[%3.3f], D_Real_Loss[%3.3f], D_Fake_Loss[%3.3f], Siamese_Loss[%3.3f]' %\n",
    "                  (batch%batch_total, batch_total, batch/float(batch_total)*100, int(np.floor(batch/float(batch_total))),\n",
    "                   gene_loss2, disc_real_loss2, disc_fake_loss2, sia_loss2))\n",
    "        \n",
    "            current_progress = elapsed / FLAGS.train_time\n",
    "          \n",
    "            if current_progress >= 1.0:\n",
    "                done = True\n",
    "            \n",
    "            # Update learning rate\n",
    "            if batch % FLAGS.learning_rate_half_life == 0:\n",
    "                lrval *= .5\n",
    "\n",
    "        if batch % FLAGS.summary_period == 0:\n",
    "            # Show progress with test features\n",
    "            feed_dict = {gene_minput: test_feature}\n",
    "            gene_output = sess.run(gene_moutput, feed_dict=feed_dict)\n",
    "            image = _summarize_progress(sess,test_feature, test_label, gene_output, batch, 'out')\n",
    "    \n",
    "            #image_summary_t = tf.summary.image(\"Train_Images\", image)\n",
    "           # image_summary = sess.run(image_summary_t)\n",
    "            \n",
    "            summary_str = sess.run(summaries)\n",
    "            summary_writer.add_summary(summary_str, batch)\n",
    "            #summary_writer.add_summary(image_summary)\n",
    "            \n",
    "        if batch % FLAGS.checkpoint_period == 0:\n",
    "            # Save checkpoint\n",
    "            _save_checkpoint(sess, batch, saver)\n",
    "\n",
    "    _save_checkpoint(sess, batch, saver)\n",
    "    print('Finished training!')\n",
    "\n",
    "def main(argv=None):\n",
    "    _train()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    tf.app.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
